/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2019-2022. All rights reserved.
 * Description: 集合通信算子插件
 * Author: lilianlin
 * Create: 2019-11-28
 */

#include "auto_tuning_plugin.h"
#include <sstream>
#include "external/ge/ge_api_types.h" // ge对内options
#include "framework/common/ge_types.h" // ge对外options
#include "hccl/hcom.h"
#include "hcom_pub.h"
#include "hccl/hccl_types.h"
#include "comm.h"

namespace hccl {
AutoTuningPlugin::AutoTuningPlugin()
    : opsKernelInfoStorePtr_(nullptr), graphOptimizerPtr_(nullptr)
{
}

AutoTuningPlugin::~AutoTuningPlugin()
{
}

ge::Status AutoTuningPlugin::Finalize()
{
    opsKernelInfoStorePtr_ = nullptr;
    graphOptimizerPtr_ = nullptr;
    return ge::SUCCESS;
}

void AutoTuningPlugin::GetOpsKernelInfoStores(map<string, OpsKernelInfoStorePtr> &opKernInfos)
{
    HCCL_INFO("get hcom kernel info store start.");

    if (opsKernelInfoStorePtr_ != nullptr) {
        opKernInfos.insert(std::make_pair(AUTOTUNE_HCCL_OPS_LIB_NAME, opsKernelInfoStorePtr_));
    } else {
        HCCL_ERROR("[Get][OpsKernelInfoStores]get hcom ops kernel info stores ptr failed for nullptr.");
    }
    HCCL_INFO("get hcom kernel info store finished.");
    return;
}

AutoTuningPlugin &AutoTuningPlugin::Instance()
{
    static AutoTuningPlugin plugin;
    return plugin;
}

void AutoTuningPlugin::GetGraphOptimizerObjs(map<string, GraphOptimizerPtr> &graphOptimizers)
{
    HCCL_INFO("get hcom graph optimizer objs start.");
#ifndef HCOM_EXECUTOR
    if (graphOptimizerPtr_ != nullptr) {
        graphOptimizers.insert(std::make_pair(HCCL_GRAPH_OPTIMIZER_NAME, graphOptimizerPtr_));
    } else {
        HCCL_ERROR("[Get][GraphOptimizerObjs]get hcom graph optimizer objs failed for nullptr.");
    }
#endif
    HCCL_INFO("get hcom graph optimizer objs end.");
    return;
}

ge::Status AutoTuningPlugin::Initialize(const std::map<string, string> &options)
{
    EXECEPTION_CATCH((opsKernelInfoStorePtr_ = std::make_shared<hccl::AutoTuningHcomOpsKernelInfoStore>()),
        return ge::INTERNAL_ERROR);
#ifndef HCOM_EXECUTOR
    EXECEPTION_CATCH((graphOptimizerPtr_ = std::make_shared<hccl::AutoTuningHcomGraphOptimizer>()),
        return ge::INTERNAL_ERROR);
#endif
    HcomSetAutoTuneMode(true);
    HCCL_INFO("Auto Tuning hcom ops plugin init success.");
    return ge::SUCCESS;
}
}
